import os
import sys

import timm
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# Image datasets and image manipulation
import torchvision
import torchvision.transforms as transforms
from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
# Image display
import matplotlib.pyplot as plt
import numpy as np
from timm.models import create_model
from torch.nn import Sequential, Linear
import argparse

from torch.utils.data import DataLoader
from torchvision import datasets
import robust_models
import torchmetrics

def test_loop(dataloader, model, device, num_classes=2, task='binary'):
    model.eval()
    # 实例化相关metrics的计算对象
    test_acc = torchmetrics.Accuracy(task=task, num_classes=num_classes).to(device)
    test_recall = torchmetrics.Recall(task=task, average='none', num_classes=num_classes).to(device)
    test_precision = torchmetrics.Precision(task=task, average='none', num_classes=num_classes).to(device)
    test_auc = torchmetrics.AUROC(task=task, average="macro", num_classes=num_classes).to(device)

    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
            # 一个batch进行计算迭代
            test_acc(pred.argmax(1), y)
            test_auc.update(pred.argmax(1), y)
            test_recall(pred.argmax(1), y)
            test_precision(pred.argmax(1), y)

    test_loss /= num_batches
    correct /= size

    # 计算一个epoch的accuray、recall、precision、AUC
    total_acc = test_acc.compute()
    total_recall = test_recall.compute()
    total_precision = test_precision.compute()
    total_auc = test_auc.compute()
    print(f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, "
          f"torch metrics acc: {(100 * total_acc):>0.1f}%\n")
    print("recall of every test dataset class: ", total_recall)
    print("precision of every test dataset class: ", total_precision)
    print("auc:", total_auc.item())

    # 清空计算对象
    test_precision.reset()
    test_acc.reset()
    test_recall.reset()
    test_auc.reset()


def test(model, test_data):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))
    classes = [
        "sensitive",
        "resistant"
    ]

    model.eval()
    x, y = test_data[0][0], test_data[0][1]
    with torch.no_grad():
        x = x.to(device)
        pred = model(x)
        predicted, actual = classes[pred[0].argmax(0)], classes[y]
        print(f'Predicted: "{predicted}", Actual: "{actual}"')


if __name__ == '__main__':
    data_dir = '../11'
    batch_size = 256
    num_workers = 8

    data_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    test_dataset = datasets.ImageFolder(os.path.join(data_dir, 'test'), data_transforms)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

    # # rvit model
    model = robust_models.rvt_tiny(pretrained=False, num_classes=2)
    checkpoint = torch.load('../models/rvit/VIT_epoch_200_valacc0.7979166666666667.pth', map_location="cuda:0" if torch.cuda.is_available() else "cpu")
    model.load_state_dict(checkpoint, strict=False)
    model.to(device)

    # classic vit model
    # model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes = 2)
    # model.head = Sequential(Linear(model.head.in_features, 2))
    # model.load_state_dict(torch.load('../models/classic_vit/VIT_epoch_200_valacc0.80375.pth', map_location="cuda:0" if torch.cuda.is_available() else "cpu"))
    # model.to(device)

    test_loop(test_loader, model, device, num_classes=2)
